package netfilter

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"sync"
	"time"

	lp "lbs-sdk/protocol"

	"lbs-sdk/utils"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcapgo"
	lru "github.com/hashicorp/golang-lru/v2"
	"gitlab.xunlei.cn/xacc/fastlogger"

	"gvisor.dev/gvisor/pkg/tcpip"
	"gvisor.dev/gvisor/pkg/tcpip/checksum"
	"gvisor.dev/gvisor/pkg/tcpip/header"
)

var (
	pcapWriter *pcapgo.Writer
	pcapFile   *os.File
	pcapMutex  sync.Mutex
)

var OpenSessionLog = false
var logSession = fastlogger.WithField("module", "session")
var SessionLog = func(fmtStr string, args ...interface{}) {
	if !OpenSessionLog {
		return
	}
	log.Printf(fmt.Sprintf("%s %s", "[session]", fmtStr), args...)
}

// 初始化pcap写入器
func initPcapWriter() error {
	pcapMutex.Lock()
	defer pcapMutex.Unlock()

	if pcapWriter != nil {
		return nil // 已经初始化
	}

	filename := fmt.Sprintf("netfilter_packets.pcap")
	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("failed to create pcap file: %v", err)
	}

	writer := pcapgo.NewWriter(file)
	err = writer.WriteFileHeader(65536, layers.LinkTypeEthernet)
	if err != nil {
		file.Close()
		return fmt.Errorf("failed to write pcap header: %v", err)
	}

	pcapFile = file
	pcapWriter = writer
	return nil
}

// 关闭pcap写入器
func closePcapWriter() {
	pcapMutex.Lock()
	defer pcapMutex.Unlock()

	if pcapFile != nil {
		pcapFile.Close()
		pcapFile = nil
		pcapWriter = nil
	}
}

// 添加以太网头部
func addEthernetHeader(ipPacket []byte, srcMAC, dstMAC net.HardwareAddr) []byte {
	// 以太网头部是14字节
	ethHeader := make([]byte, 14)

	// 目标MAC地址 (6字节)
	copy(ethHeader[0:6], dstMAC)

	// 源MAC地址 (6字节)
	copy(ethHeader[6:12], srcMAC)

	// EtherType: IPv4 = 0x0800 (2字节，大端序)
	binary.BigEndian.PutUint16(ethHeader[12:14], 0x0800)

	// 合并以太网头部和IP包
	result := make([]byte, len(ethHeader)+len(ipPacket))
	copy(result[0:len(ethHeader)], ethHeader)
	copy(result[len(ethHeader):], ipPacket)

	return result
}

// 写入数据包到pcap文件
func writeToPcap(packet []byte) error {
	pcapMutex.Lock()
	defer pcapMutex.Unlock()

	if pcapWriter == nil {
		if err := initPcapWriter(); err != nil {
			return err
		}
	}

	// 创建捕获信息
	ci := gopacket.CaptureInfo{
		Timestamp:     time.Now(),
		CaptureLength: len(packet),
		Length:        len(packet),
	}

	return pcapWriter.WritePacket(ci, packet)
}

type NfIpPacketOptions struct {
	IpFamily          uint16 // 暂时不使用：默认 AF_INET = 2
	IpHeaderSize      int    // 暂时不使用：默认 20
	CompartmentId     uint32
	InterfaceIndex    uint32
	SubInterfaceIndex uint32
	Flags             uint32
}

// 接收包
type NetFilterRecvIpPacket func(pkg []byte, version uint8) Verdict

// 发送包
type NetFilterSendIpPacket func(pkg []byte, options NfIpPacketOptions) error

type FilterFnList struct {
	Idx          int
	Size         int
	FilterFnList []NetFilterRecvIpPacket
}

func CreateFilterFnList(size int) *FilterFnList {
	return &FilterFnList{
		Idx:          0,
		Size:         size,
		FilterFnList: make([]NetFilterRecvIpPacket, size),
	}
}

func (flst *FilterFnList) AddFilterFnList(fn NetFilterRecvIpPacket) Verdict {
	if flst.Idx >= flst.Size {
		return NF_ACCEPT
	}
	flst.FilterFnList[flst.Idx] = fn
	flst.Idx++
	return NF_ACCEPT
}

func (flst *FilterFnList) Filter(pkg []byte, ipVersion uint8) Verdict {
	for i := 0; i < flst.Idx; i++ {
		fn := flst.FilterFnList[i]
		if fn == nil {
			continue
		}
		switch fn(pkg, ipVersion) {
		case NF_ACCEPT:
			continue
		case NF_DROP:
			return NF_DROP
		case NF_STOLEN:
			return NF_STOLEN
		default:
			return NF_DROP
		}
	}
	return NF_ACCEPT
}

type OutPut func(pkg []byte)

// netfilter client driver
type ClientDriverNf struct {
	Sessions  *lru.Cache[Tun5, IpSession] // session cache
	OutPut    OutPut
	VpnRemote []Tun5 // 远程VPN地址列表
	Ctx       context.Context
	Cancel    context.CancelFunc

	FakeIpv4 tcpip.Address
	FakeIpv6 tcpip.Address

	Writer io.Writer
}

var DefaultFilterFnListCache = CreateFilterFnList(16)
var DefaultClientDriverNf = NewClientDriverNf()

func NewClientDriverNf() *ClientDriverNf {
	ctx, cancel := context.WithCancel(context.Background())
	sessions, err := lru.New[Tun5, IpSession](2048)
	if err != nil {
		panic(fmt.Sprintf("create session cache failed: %v", err))
	}
	return &ClientDriverNf{
		Sessions: sessions,
		Ctx:      ctx,
		Cancel:   cancel,
	}
}

type Tun5 struct {
	Src      string
	SrcPort  uint16
	Dst      string
	DstPort  uint16
	Protocol uint8
}

var ProxyList = []string{"192.168.182.63", "8.217.11.128"}
var ProxyAll = false

func IPv4ToIPv6(ipv4 net.IP) net.IP {
	ipv4 = ipv4.To4()
	if ipv4 == nil {
		return nil
	}
	return net.IP(append([]byte{0xF, 0xC, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, ipv4...))
}

func (c *ClientDriverNf) SetFakeIP(strIp string) {
	ip := net.ParseIP(strIp)
	if ip == nil {
		panic("invalid fake IP")
	}
	c.FakeIpv4 = tcpip.AddrFrom4Slice(ip.To4())
	c.FakeIpv6 = tcpip.AddrFrom16([16]byte(IPv4ToIPv6(ip)))
}

func (c *ClientDriverNf) SetWriter(w io.Writer) {
	c.Writer = w
}

func checksumUpdate2ByteAlignedAddress(xsum uint16, old, new tcpip.Address) uint16 {
	const uint16Bytes = 2

	if old.BitLen() != new.BitLen() {
		panic(fmt.Sprintf("buffer lengths are different; old = %d, new = %d", old.BitLen()/8, new.BitLen()/8))
	}

	if oldBytes := old.BitLen() % 16; oldBytes != 0 {
		panic(fmt.Sprintf("buffer has an odd number of bytes; got = %d", oldBytes))
	}

	oldAddr := old.AsSlice()
	newAddr := new.AsSlice()

	// As per RFC 1071 page 4,
	//	(4)  Incremental Update
	//
	//        ...
	//
	//        To update the checksum, simply add the differences of the
	//        sixteen bit integers that have been changed.  To see why this
	//        works, observe that every 16-bit integer has an additive inverse
	//        and that addition is associative.  From this it follows that
	//        given the original value m, the new value m', and the old
	//        checksum C, the new checksum C' is:
	//
	//                C' = C + (-m) + m' = C + (m' - m)
	for len(oldAddr) != 0 {
		// Convert the 2 byte sequences to uint16 values then apply the increment
		// update.
		xsum = checksumUpdate2ByteAlignedUint16(xsum, (uint16(oldAddr[0])<<8)+uint16(oldAddr[1]), (uint16(newAddr[0])<<8)+uint16(newAddr[1]))
		oldAddr = oldAddr[uint16Bytes:]
		newAddr = newAddr[uint16Bytes:]
	}

	return xsum
}
func checksumUpdate2ByteAlignedUint16(xsum, old, new uint16) uint16 {
	// As per RFC 1071 page 4,
	//	(4)  Incremental Update
	//
	//        ...
	//
	//        To update the checksum, simply add the differences of the
	//        sixteen bit integers that have been changed.  To see why this
	//        works, observe that every 16-bit integer has an additive inverse
	//        and that addition is associative.  From this it follows that
	//        given the original value m, the new value m', and the old
	//        checksum C, the new checksum C' is:
	//
	//                C' = C + (-m) + m' = C + (m' - m)
	return checksum.Combine(xsum, checksum.Combine(new, ^old))
}

func (c *ClientDriverNf) IpPkgFromNetfilter(ipPkg []byte, ipVersion uint8, options NfIpPacketOptions) (verdict Verdict) {
	// 这个方法处理的包一定不能再交给协议栈处理
	// 因为方法会修改数据包内容，可能导致环路
	utils.PcapDump(ipPkg)
	var protocol tcpip.TransportProtocolNumber
	var oldSrcIp, oldDstIp, newSrcIp, newDstIp tcpip.Address
	var l4Pkg []byte
	var fakeIp tcpip.Address
	switch ipVersion {
	case header.IPv4Version:
		ipv4Hdr := header.IPv4(ipPkg)
		protocol = ipv4Hdr.TransportProtocol()
		l4Pkg = ipv4Hdr.Payload()
		oldSrcIp = ipv4Hdr.SourceAddress()
		oldDstIp = ipv4Hdr.DestinationAddress()
		if ipv4Hdr.FragmentOffset() != 0 {
			SessionLog("fragmented ipv4 packet")
			return
		}
		ipHeaderLen := int(ipv4Hdr.HeaderLength())
		// 重新计算ipv4头部校验和
		ipv4Hdr.SetSourceAddress(c.FakeIpv4)
		ipv4Hdr.SetChecksum(0)
		csm := lp.CalculateChecksum(ipPkg[:ipHeaderLen], 0)
		ipv4Hdr.SetChecksum(csm)
		fakeIp = c.FakeIpv4
		newSrcIp, newDstIp = c.FakeIpv4, oldDstIp
	case header.IPv6Version:
		ipv6Hdr := header.IPv6(ipPkg)
		protocol = ipv6Hdr.TransportProtocol()
		l4Pkg = ipv6Hdr.Payload()
		oldSrcIp = ipv6Hdr.SourceAddress()
		oldDstIp = ipv6Hdr.DestinationAddress()
		fakeIp = c.FakeIpv6
		if protocol == header.IPv6FragmentHeader {
			expHdr := header.IPv6Fragment(l4Pkg)
			if !expHdr.IsValid() {
				SessionLog("invalid ipv6 fragment header")
				return
			}
			if expHdr.More() && expHdr.FragmentOffset() == 0 {
				protocol = tcpip.TransportProtocolNumber(expHdr.NextHeader())
				l4Pkg = expHdr.Payload()
			} else {
				SessionLog("fragmented ipv6 packet")
				return
			}
		}
		ipv6Hdr.SetSourceAddress(c.FakeIpv6)
		newSrcIp, newDstIp = c.FakeIpv6, oldDstIp
	default:
		SessionLog("unknown IP version: %d", ipVersion)
		return
	}
	// 一个localip对应一个nat table
	tb := GetNatTable(ipVersion)
	switch protocol {
	case header.TCPProtocolNumber:
		if len(l4Pkg) < header.TCPMinimumSize {
			SessionLog("tcp header too short: %d bytes", len(l4Pkg))
			return
		}
		tcp := header.TCP(l4Pkg)
		oldSrcPort := tcp.SourcePort()
		peer, ok := tb.AddSession(options, fakeIp, oldSrcIp, oldSrcPort, oldDstIp, tcp.DestinationPort(), protocol)
		if !ok {
			SessionLog("tcp session full, drop it")
			return
		}
		SessionLog("tcp %v %s:%d -> %s:%d natto %s:%d\n",
			ok, oldSrcIp.String(), oldSrcPort, oldDstIp.String(), tcp.DestinationPort(), peer.FakeIP, peer.FakePort)
		tcp.SetSourcePort(peer.FakePort)
		tcp.SetChecksum(0)
		csm := lp.RecalculateTransportChecksum(l4Pkg, ipVersion, newSrcIp, newDstIp, uint8(protocol))
		tcp.SetChecksum(csm)
		// tcp.SetSourcePortWithChecksumUpdate(peer.FakePort)
		// tcp.UpdateChecksumPseudoHeaderAddress(oldSrcIp, peer.FakeIP, true)
	case header.UDPProtocolNumber:
		if len(l4Pkg) < header.UDPMinimumSize {
			SessionLog("udp header too short: %d bytes", len(l4Pkg))
			return
		}
		udp := header.UDP(l4Pkg)
		oldSrcPort := udp.SourcePort()
		peer, ok := tb.AddSession(options, fakeIp, oldSrcIp, oldSrcPort, oldDstIp, udp.DestinationPort(), protocol)
		if !ok {
			SessionLog("udp session full, drop it")
			return
		}
		SessionLog("udp %v %s:%d -> %s:%d natto %s:%d\n",
			ok, oldSrcIp.String(), oldSrcPort, oldDstIp.String(), udp.DestinationPort(), peer.FakeIP, peer.FakePort)
		udp.SetChecksum(0)
		csm := lp.RecalculateTransportChecksum(l4Pkg, ipVersion, newSrcIp, newDstIp, uint8(protocol))
		udp.SetChecksum(csm)
		// udp.SetSourcePortWithChecksumUpdate(peer.FakePort)
		// udp.UpdateChecksumPseudoHeaderAddress(oldSrcIp, peer.FakeIP, true)
	case header.ICMPv4ProtocolNumber:
		if len(l4Pkg) < header.ICMPv4MinimumSize {
			SessionLog("ICMP header too short: %d bytes", len(l4Pkg))
			return
		}
		icmp := header.ICMPv4(l4Pkg)
		oldIdent := icmp.Ident()
		peer, ok := tb.AddSession(options, fakeIp, oldSrcIp, oldIdent, oldDstIp, oldIdent, protocol)
		if !ok {
			SessionLog("icmp session full, drop it")
			return
		}
		SessionLog("icmp %v %s -> %s, ident=%d natto %s:%d\n",
			ok, oldSrcIp.String(), oldDstIp.String(), oldIdent, peer.FakeIP, peer.FakePort)
		icmp.SetChecksum(0)
		icmp.SetIdent(peer.FakePort)
		csm := lp.RecalculateIcmpChecksum(l4Pkg, ipVersion, newSrcIp, newDstIp)
		icmp.SetChecksum(csm)
		// icmp.SetIdentWithChecksumUpdate(peer.FakePort)
	case header.ICMPv6ProtocolNumber:
		if len(l4Pkg) < header.ICMPv6MinimumSize {
			SessionLog("ICMPv6 header too short: %d bytes", len(l4Pkg))
			return
		}
		icmpv6 := header.ICMPv6(l4Pkg)
		oldIdent := icmpv6.Ident()
		peer, ok := tb.AddSession(options, fakeIp, oldSrcIp, oldIdent, oldDstIp, oldIdent, protocol)
		if !ok {
			SessionLog("icmpv6 session full, drop it")
			return
		}
		// ipv6Hdr := header.IPv6(ipPkg)
		// newCsm := checksumUpdate2ByteAlignedAddress(^icmpv6.Checksum(), oldSrcIp, ipv6Hdr.SourceAddress())
		// newCsm = ^checksumUpdate2ByteAlignedUint16(newCsm, oldIdent, peer.FakePort)
		icmpv6.SetIdent(peer.FakePort)
		icmpv6.SetChecksum(0)
		csm := lp.RecalculateTransportChecksum(l4Pkg, ipVersion, newSrcIp, newDstIp, uint8(protocol))
		icmpv6.SetChecksum(csm)
		// icmpv6.SetChecksum(newCsm)
		SessionLog("icmpv6 %v %s -> %s, ident=%d natto %s:%d\n", ok, oldSrcIp.String(), oldDstIp.String(), oldIdent, peer.FakeIP, peer.FakePort)
	default:
		SessionLog("unknown protocol: %d", protocol)
		return
	}
	var fragments [][]byte
	var err error
	switch ipVersion {
	case header.IPv4Version:
		fragments = lp.FragmentPayload(ipPkg, utils.MTU)
	case header.IPv6Version:
		fragments, err = lp.FragmentPayloadIpv6(ipPkg, utils.MTU)
		if err != nil {
			SessionLog("ipv6 fragment error: %v", err)
			return
		}
	default:
		panic("unknown ip version")
	}
	for _, frag := range fragments {
		utils.PcapDump(frag)
		if c.Writer != nil {
			c.Writer.Write(frag)
		}
	}
	return
}

func (c *ClientDriverNf) IpPkgToNetfilter(ipPkg []byte) error {
	utils.PcapDump(ipPkg)
	ipVersion := header.IPVersion(ipPkg)

	var session *IpSession
	var fragments [][]byte
	var oldPort uint16
	// var getPort bool
	switch ipVersion {
	case header.IPv4Version:
		ipv4Hdr := header.IPv4(ipPkg)
		if ipv4Hdr.FragmentOffset() != 0 || ipv4Hdr.More() {
			// 处理分片
			complete, reassembled, err := lp.DefaultIPv4Reassembler.AddFragment(ipPkg)
			if err != nil {
				return err
			}
			if !complete {
				// 尚未重组完成，等待更多分片
				fmt.Printf("waiting more fragments for ipv4\n")
				return nil
			}
			fragments = [][]byte{reassembled}
			// getPort, oldPort, fragments = lp.DefaultIpFragmentManager.ProcessIpFragment(header.IPv4Version, ipPkg)
			// if !getPort {
			// 	return nil
			// }
		} else {
			fragments = [][]byte{ipPkg}
		}
		for _, frag := range fragments {
			session = ProcessIpv4(frag, oldPort)
			if session == nil {
				continue
			}
			session.Options.IpFamily = 2
			session.Options.IpHeaderSize = 20
			utils.PcapDump(frag)
			NetfilterSendIpPkg(frag, session.Options)
		}
		return nil
	case header.IPv6Version:
		panic("not implement ipv6 yet")
	default:
		panic("unknown ip version")
	}
}

const USE_FIXED_FAKE_IP = true

var FakeIp net.IP

// 生成10.200.1.1/13网段的随机IP
func generateFakeIP() net.IP {
	if USE_FIXED_FAKE_IP {
		// 返回固定IP: 10.200.1.121
		return FakeIp
	}
	// 10.200.0.0/13 = 10.200.0.0 - 10.207.255.255
	// 基础IP: 10.200.0.0
	baseIP := net.IPv4(10, 200, 0, 0)

	// 生成随机数
	var randomBytes [4]byte
	rand.Read(randomBytes[:])

	// 13位掩码意味着19位主机位 (32-13=19)
	// 掩码: 11111111.11111000.00000000.00000000
	hostBits := binary.BigEndian.Uint32(randomBytes[:]) & 0x0007FFFF // 保留低19位

	// 确保不是网络地址或广播地址
	if hostBits == 0 {
		hostBits = 1
	}
	if hostBits == 0x0007FFFF {
		hostBits = 0x0007FFFE
	}

	// 将随机主机位应用到基础IP
	baseIPInt := binary.BigEndian.Uint32(baseIP.To4())
	fakeIPInt := baseIPInt | hostBits

	fakeIP := make(net.IP, 4)
	binary.BigEndian.PutUint32(fakeIP, fakeIPInt)

	return fakeIP
}

// 从五元组解析
func parseTuple5(packet gopacket.Packet) (*Tun5, error) {
	var tun5 Tun5

	// 解析IP层
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return nil, fmt.Errorf("no IPv4 layer found")
	}

	ip, ok := ipLayer.(*layers.IPv4)
	if !ok {
		return nil, fmt.Errorf("failed to parse IPv4 layer")
	}

	tun5.Src = ip.SrcIP.String()
	tun5.Dst = ip.DstIP.String()
	tun5.Protocol = uint8(ip.Protocol)

	// 解析传输层端口
	switch ip.Protocol {
	case layers.IPProtocolTCP:
		if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
			tcp, _ := tcpLayer.(*layers.TCP)
			tun5.SrcPort = uint16(tcp.SrcPort)
			tun5.DstPort = uint16(tcp.DstPort)
		}
	case layers.IPProtocolUDP:
		if udpLayer := packet.Layer(layers.LayerTypeUDP); udpLayer != nil {
			udp, _ := udpLayer.(*layers.UDP)
			tun5.SrcPort = uint16(udp.SrcPort)
			tun5.DstPort = uint16(udp.DstPort)
		}
	}

	return &tun5, nil
}
